# query_vector_store.py
import os

import redis

from services.vector_store import VectorStoreService

REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))

r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0, decode_responses=True)
def get_all():
    # 初始化向量存储服务
    vector_service = VectorStoreService()

    if not vector_service.vector_store:
        print("向量存储未初始化")
        return

    print("=== 向量库查询工具 ===")

    # # 显示文档总数
    # total = vector_service.get_document_count()
    # print(f"文档总数: {total}")

    # 显示前5个文档的详细信息
    print("\n前5个文档:")
    documents = vector_service.list_all_documents()
    for i, doc in enumerate(documents, 1):
        print(f"\n文档 {i}:")
        print(f"  ID: {doc['id']}")
        print(f"  内容预览: {doc['content']}...")
        print(f"  元数据: {doc['metadata']}")


def del_doc():
    # 初始化向量存储服务
    vector_service = VectorStoreService()

    if not vector_service.vector_store:
        print("向量存储未初始化")
        return
    ids = ["ad410769-80b5-48c5-b9ae-605f25068d86","a1d04a31-6b23-4bee-a590-75f3c4ce720a"]
    vector_service.delete_documents_by_ids(ids)

def get_all_rds():
    doc_ids = r.smembers("all_documents")
    doc_ids_list = list(doc_ids)
    print(doc_ids_list)

def get_msg():
    vector_service = VectorStoreService()
    docs = vector_service.similarity_search("没有收入能不能办", k=1)
    print(docs[0])

def main():
    # get_msg()
    get_all()
    # del_doc()
    # get_all_rds()

if __name__ == "__main__":
    main()
